// This file is part of Notepad4.
// See License.txt for details about distribution and modification.
//! Lexer for Kotlin.

#include <cassert>
#include <cstring>

#include <string>
#include <string_view>
#include <vector>

#include "ILexer.h"
#include "Scintilla.h"
#include "SciLexer.h"

#include "WordList.h"
#include "LexAccessor.h"
#include "Accessor.h"
#include "StyleContext.h"
#include "CharacterSet.h"
#include "StringUtils.h"
#include "LexerModule.h"
#include "LexerUtils.h"

using namespace Lexilla;

namespace {

struct EscapeSequence {
	int outerState = SCE_KOTLIN_DEFAULT;
	int digitsLeft = 0;

	// highlight any character as escape sequence.
	bool resetEscapeState(int state, int chNext) noexcept {
		if (IsEOLChar(chNext)) {
			return false;
		}
		outerState = state;
		digitsLeft = (chNext == 'u') ? 5 : 1;
		return true;
	}
	bool atEscapeEnd(int ch) noexcept {
		--digitsLeft;
		return digitsLeft <= 0 || !IsHexDigit(ch);
	}
};

enum {
	KotlinLineStateMaskLineComment = 1, // line comment
	KotlinLineStateMaskImport = 1 << 1, // import
};

//KeywordIndex++Autogenerated -- start of section automatically generated
enum {
	KeywordIndex_Keyword = 0,
	KeywordIndex_JavaClass = 1,
	KeywordIndex_Class = 2,
	KeywordIndex_JavaInterface = 3,
	KeywordIndex_Interface = 4,
	KeywordIndex_Enumeration = 5,
	MaxKeywordSize = 32,
};
//KeywordIndex--Autogenerated -- end of section automatically generated

enum class KeywordType {
	None = SCE_KOTLIN_DEFAULT,
	Label = SCE_KOTLIN_LABEL,
	ThisSuper = SCE_KOTLIN_IDENTIFIER,
	Annotation = SCE_KOTLIN_ANNOTATION,
	Class = SCE_KOTLIN_CLASS,
	Interface = SCE_KOTLIN_INTERFACE,
	Enum = SCE_KOTLIN_ENUM,
	Function = SCE_KOTLIN_FUNCTION_DEFINITION,
	Return = 0x40,
};

static_assert(DefaultNestedStateBaseStyle + 1 == SCE_KOTLIN_STRING);
static_assert(DefaultNestedStateBaseStyle + 2 == SCE_KOTLIN_RAWSTRING);

constexpr bool IsSpaceEquiv(int state) noexcept {
	return state <= SCE_KOTLIN_TASKMARKER;
}

void ColouriseKotlinDoc(Sci_PositionU startPos, Sci_Position lengthDoc, int initStyle, LexerWordList keywordLists, Accessor &styler) {
	int lineStateLineType = 0;
	int commentLevel = 0;	// nested block comment level

	KeywordType kwType = KeywordType::None;
	int chBeforeIdentifier = 0;

	std::vector<int> nestedState; // string interpolation "${}"

	int visibleChars = 0;
	int chBefore = 0;
	int visibleCharsBefore = 0;
	int chPrevNonWhite = 0;
	EscapeSequence escSeq;

	StyleContext sc(startPos, lengthDoc, initStyle, styler);
	if (sc.currentLine > 0) {
		int lineState = styler.GetLineState(sc.currentLine - 1);
		/*
		2: lineStateLineType
		6: commentLevel
		3: nestedState count
		3*4: nestedState
		*/
		commentLevel = (lineState >> 2) & 0x3f;
		lineState >>= 8;
		if (lineState) {
			UnpackLineState(lineState, nestedState);
		}
	} else if (startPos == 0 && sc.Match('#', '!')) {
		// Shell Shebang at beginning of file
		sc.SetState(SCE_KOTLIN_COMMENTLINE);
		sc.Forward();
		lineStateLineType = KotlinLineStateMaskLineComment;
	}

	while (sc.More()) {
		switch (sc.state) {
		case SCE_KOTLIN_OPERATOR:
		case SCE_KOTLIN_OPERATOR2:
			sc.SetState(SCE_KOTLIN_DEFAULT);
			break;

		case SCE_KOTLIN_NUMBER:
			if (!IsDecimalNumber(sc.chPrev, sc.ch, sc.chNext)) {
				sc.SetState(SCE_KOTLIN_DEFAULT);
			}
			break;

		case SCE_KOTLIN_LABEL:
		case SCE_KOTLIN_IDENTIFIER:
		case SCE_KOTLIN_ANNOTATION:
		case SCE_KOTLIN_CLASS:
			if (!IsIdentifierCharEx(sc.ch)) {
				if (sc.state == SCE_KOTLIN_ANNOTATION) {
					if (sc.ch == '.' || sc.ch == ':') {
						sc.SetState(SCE_KOTLIN_OPERATOR);
						sc.ForwardSetState(SCE_KOTLIN_ANNOTATION);
						continue;
					}
				} else if (sc.state == SCE_KOTLIN_IDENTIFIER) {
					char s[MaxKeywordSize];
					sc.GetCurrent(s, sizeof(s));
					if (keywordLists[KeywordIndex_Keyword].InList(s)) {
						sc.ChangeState(SCE_KOTLIN_WORD);
						if (escSeq.outerState != SCE_KOTLIN_DEFAULT) {
							kwType = KeywordType::None;
						} else if (StrEqual(s, "import")) {
							if (visibleChars == sc.LengthCurrent()) {
								lineStateLineType = KotlinLineStateMaskImport;
							}
						} else if (sc.ch == '@' && StrEqualsAny(s, "break", "continue", "return")) {
							kwType = KeywordType::Label;
						} else if (sc.ch == '@' && StrEqualsAny(s, "this", "super")) {
							kwType = KeywordType::ThisSuper;
						} else if (StrEqualsAny(s, "class", "typealias", "throw", "is", "as")) {
							if (kwType != KeywordType::Annotation && kwType != KeywordType::Enum) {
								kwType = KeywordType::Class;
							}
						} else if (StrEqual(s, "enum")) {
							kwType = KeywordType::Enum;
						} else if (StrEqual(s, "fun")) {
							kwType = KeywordType::Function;
						} else if (StrEqual(s, "annotation")) {
							kwType = KeywordType::Annotation;
						} else if (StrEqual(s, "interface")) {
							kwType = KeywordType::Interface;
						} else if (StrEqualsAny(s, "return", "else")) {
							kwType = KeywordType::Return;
						}
						if (kwType > KeywordType::ThisSuper && kwType < KeywordType::Return) {
							const int chNext = sc.GetDocNextChar();
							if (chNext != '?' && !IsIdentifierStartEx(chNext)) {
								kwType = KeywordType::None;
							}
						}
					} else if (escSeq.outerState == SCE_KOTLIN_DEFAULT && sc.ch == '@') {
						sc.ChangeState(SCE_KOTLIN_LABEL);
						sc.Forward();
					} else if (keywordLists[KeywordIndex_JavaClass].InList(s) || keywordLists[KeywordIndex_Class].InList(s)) {
						sc.ChangeState(SCE_KOTLIN_CLASS);
					} else if (keywordLists[KeywordIndex_JavaInterface].InList(s) || keywordLists[KeywordIndex_Interface].InList(s)) {
						sc.ChangeState(SCE_KOTLIN_INTERFACE);
					} else if (keywordLists[KeywordIndex_Enumeration].InList(s)) {
						sc.ChangeState(SCE_KOTLIN_ENUM);
					} else if (escSeq.outerState == SCE_KOTLIN_DEFAULT) {
						if (sc.ch == '.') {
							if (kwType == KeywordType::Function) {
								// extension function
								sc.ChangeState(SCE_KOTLIN_CLASS);
							}
						} else if (kwType > KeywordType::None && kwType < KeywordType::Return) {
							sc.ChangeState(static_cast<int>(kwType));
						} else {
							const int chNext = sc.GetDocNextChar(sc.ch == '?');
							if (chNext == '(') {
								// type function()
								// type[] function()
								// type<type> function()
								if (kwType != KeywordType::Return && (IsIdentifierCharEx(chBefore) || chBefore == ']')) {
									sc.ChangeState(SCE_KOTLIN_FUNCTION_DEFINITION);
								} else {
									sc.ChangeState(SCE_KOTLIN_FUNCTION);
								}
							} else if (sc.Match(':', ':')
								|| (chBeforeIdentifier == '<' && (chNext == '>' || chNext == '<'))
								|| (chBeforeIdentifier == ':' && (sc.ch == '?' && chNext != '.'))) {
								// type::class
								// type<type>
								// type<type?>
								// type<type<type>>
								// type<type, type>
								// class type: type, interface {}
								// identifier: type?
								sc.ChangeState(SCE_KOTLIN_CLASS);
							}
						}
					}
					if (sc.state != SCE_KOTLIN_WORD && sc.ch != '.') {
						kwType = KeywordType::None;
					}
					if (escSeq.outerState != SCE_KOTLIN_DEFAULT) {
						sc.SetState(escSeq.outerState);
						continue;
					}
				}
				sc.SetState(SCE_KOTLIN_DEFAULT);
			}
			break;

		case SCE_KOTLIN_COMMENTLINE:
		case SCE_KOTLIN_COMMENTLINEDOC:
			if (sc.atLineStart) {
				sc.SetState(SCE_KOTLIN_DEFAULT);
			} else {
				HighlightTaskMarker(sc, visibleChars, visibleCharsBefore, SCE_KOTLIN_TASKMARKER);
			}
			break;

		case SCE_KOTLIN_COMMENTBLOCK:
		case SCE_KOTLIN_COMMENTBLOCKDOC:
			if (sc.state == SCE_KOTLIN_COMMENTBLOCKDOC && sc.ch == '@' && IsLowerCase(sc.chNext) && IsCommentTagPrev(sc.chPrev)) {
				sc.SetState(SCE_KOTLIN_COMMENTDOCWORD);
			} else if (sc.Match('*', '/')) {
				sc.Forward();
				--commentLevel;
				if (commentLevel == 0) {
					sc.ForwardSetState(SCE_KOTLIN_DEFAULT);
				}
			} else if (sc.Match('/', '*')) {
				sc.Forward();
				++commentLevel;
			} else if (HighlightTaskMarker(sc, visibleChars, visibleCharsBefore, SCE_KOTLIN_TASKMARKER)) {
				continue;
			}
			break;

		case SCE_KOTLIN_COMMENTDOCWORD:
			if (!IsLowerCase(sc.ch)) {
				sc.SetState(SCE_KOTLIN_COMMENTBLOCKDOC);
				continue;
			}
			break;

		case SCE_KOTLIN_CHARACTER:
		case SCE_KOTLIN_STRING:
		case SCE_KOTLIN_RAWSTRING:
			if (sc.atLineStart && sc.state != SCE_KOTLIN_RAWSTRING) {
				sc.SetState(SCE_KOTLIN_DEFAULT);
			} else if (sc.ch == '\\' && sc.state != SCE_KOTLIN_RAWSTRING) {
				if (escSeq.resetEscapeState(sc.state, sc.chNext)) {
					sc.SetState(SCE_KOTLIN_ESCAPECHAR);
					sc.Forward();
				}
			} else if (sc.ch == '\'' && sc.state == SCE_KOTLIN_CHARACTER) {
				sc.ForwardSetState(SCE_KOTLIN_DEFAULT);
			} else if (sc.state != SCE_KOTLIN_CHARACTER) {
				if (sc.ch == '$') {
					if (sc.chNext == '{' || IsIdentifierStartEx(sc.chNext)) {
						escSeq.outerState = sc.state;
						sc.SetState(SCE_KOTLIN_OPERATOR2);
						sc.Forward();
						if (sc.ch == '{') {
							nestedState.push_back(escSeq.outerState);
						} else {
							sc.SetState(SCE_KOTLIN_IDENTIFIER);
						}
					}
				} else if (sc.ch == '\"' && (sc.state == SCE_KOTLIN_STRING || sc.MatchNext('"', '"'))) {
					if (sc.state == SCE_KOTLIN_RAWSTRING) {
						// quotes except last three are string content
						while (sc.chNext == '\"') {
							sc.Forward();
						}
					}
					sc.ForwardSetState(SCE_KOTLIN_DEFAULT);
				}
			}
			break;

		case SCE_KOTLIN_ESCAPECHAR:
			if (escSeq.atEscapeEnd(sc.ch)) {
				sc.SetState(escSeq.outerState);
				continue;
			}
			break;

		case SCE_KOTLIN_BACKTICKS:
			if (sc.atLineStart) {
				sc.SetState(SCE_KOTLIN_DEFAULT);
			} else if (sc.ch == '`') {
				sc.ForwardSetState(SCE_KOTLIN_DEFAULT);
			}
			break;
		}

		if (sc.state == SCE_KOTLIN_DEFAULT) {
			if (sc.ch == '/' && (sc.chNext == '/' || sc.chNext == '*')) {
				visibleCharsBefore = visibleChars;
				const int chNext = sc.chNext;
				sc.SetState((chNext == '/') ? SCE_KOTLIN_COMMENTLINE : SCE_KOTLIN_COMMENTBLOCK);
				sc.Forward(2);
				if (sc.ch == '!' || (sc.ch == chNext && sc.chNext != chNext)) {
					static_assert(SCE_KOTLIN_COMMENTLINEDOC - SCE_KOTLIN_COMMENTLINE == SCE_KOTLIN_COMMENTBLOCKDOC - SCE_KOTLIN_COMMENTBLOCK);
					sc.ChangeState(sc.state + SCE_KOTLIN_COMMENTLINEDOC - SCE_KOTLIN_COMMENTLINE);
				}
				if (chNext == '/') {
					if (visibleChars == 0) {
						lineStateLineType = KotlinLineStateMaskLineComment;
					}
				} else {
					commentLevel = 1;
				}
				continue;
			}
			if (sc.ch == '\"') {
				sc.SetState(SCE_KOTLIN_STRING);
				if (sc.MatchNext('"', '"')) {
					sc.ChangeState(SCE_KOTLIN_RAWSTRING);
					sc.Forward(2);
				}
			} else if (sc.ch == '\'') {
				sc.SetState(SCE_KOTLIN_CHARACTER);
			} else if (IsNumberStartEx(sc.chPrev, sc.ch, sc.chNext)) {
				sc.SetState(SCE_KOTLIN_NUMBER);
			} else if (sc.ch == '@' && IsIdentifierStartEx(sc.chNext)) {
				int state;
				if (kwType == KeywordType::ThisSuper || chPrevNonWhite == '>') {
					// super<type>@class
					state = SCE_KOTLIN_CLASS;
					sc.SetState(SCE_KOTLIN_OPERATOR);
					sc.Forward();
				} else {
					state = (kwType == KeywordType::Label) ? SCE_KOTLIN_LABEL : SCE_KOTLIN_ANNOTATION;
				}
				sc.SetState(state);
				kwType = KeywordType::None;
			} else if (sc.ch == '`') {
				sc.SetState(SCE_KOTLIN_BACKTICKS);
			} else if (IsIdentifierStartEx(sc.ch)) {
				escSeq.outerState = SCE_KOTLIN_DEFAULT;
				chBefore = chPrevNonWhite;
				if (chPrevNonWhite != '.') {
					chBeforeIdentifier = chPrevNonWhite;
				}
				sc.SetState(SCE_KOTLIN_IDENTIFIER);
			} else if (IsAGraphic(sc.ch)) {
				sc.SetState(SCE_KOTLIN_OPERATOR);
				if (!nestedState.empty()) {
					sc.ChangeState(SCE_KOTLIN_OPERATOR2);
					if (sc.ch == '{') {
						nestedState.push_back(SCE_KOTLIN_DEFAULT);
					} else if (sc.ch == '}') {
						const int outerState = TakeAndPop(nestedState);
						sc.ForwardSetState(outerState);
						continue;
					}
				}
			}
		}

		if (!isspacechar(sc.ch)) {
			visibleChars++;
			if (!IsSpaceEquiv(sc.state)) {
				chPrevNonWhite = sc.ch;
			}
		}
		if (sc.atLineEnd) {
			int lineState = (commentLevel << 2) | lineStateLineType;
			if (!nestedState.empty()) {
				lineState |= PackLineState(nestedState) << 8;
			}
			styler.SetLineState(sc.currentLine, lineState);
			lineStateLineType = 0;
			visibleChars = 0;
			visibleCharsBefore = 0;
			kwType = KeywordType::None;
		}
		sc.Forward();
	}

	sc.Complete();
}

struct FoldLineState {
	int lineComment;
	int packageImport;
	constexpr explicit FoldLineState(int lineState) noexcept:
		lineComment(lineState & KotlinLineStateMaskLineComment),
		packageImport((lineState >> 1) & 1) {
	}
};

void FoldKotlinDoc(Sci_PositionU startPos, Sci_Position lengthDoc, int initStyle, LexerWordList /*keywordLists*/, Accessor &styler) {
	const Sci_PositionU endPos = startPos + lengthDoc;
	Sci_Line lineCurrent = styler.GetLine(startPos);
	FoldLineState foldPrev(0);
	int levelCurrent = SC_FOLDLEVELBASE;
	if (lineCurrent > 0) {
		levelCurrent = styler.LevelAt(lineCurrent - 1) >> 16;
		foldPrev = FoldLineState(styler.GetLineState(lineCurrent - 1));
		const Sci_PositionU bracePos = CheckBraceOnNextLine(styler, lineCurrent - 1, SCE_KOTLIN_OPERATOR, SCE_KOTLIN_TASKMARKER);
		if (bracePos) {
			startPos = bracePos + 1; // skip the brace
		}
	}

	int levelNext = levelCurrent;
	FoldLineState foldCurrent(styler.GetLineState(lineCurrent));
	Sci_PositionU lineStartNext = styler.LineStart(lineCurrent + 1);

	char chNext = styler[startPos];
	int styleNext = styler.StyleIndexAt(startPos);
	int style = initStyle;
	int visibleChars = 0;

	while (startPos < endPos) {
		const char ch = chNext;
		const int stylePrev = style;
		style = styleNext;
		chNext = styler[++startPos];
		styleNext = styler.StyleIndexAt(startPos);

		switch (style) {
		case SCE_KOTLIN_COMMENTBLOCK:
		case SCE_KOTLIN_COMMENTBLOCKDOC: {
			const int level = (ch == '/' && chNext == '*') ? 1 : ((ch == '*' && chNext == '/') ? -1 : 0);
			if (level != 0) {
				levelNext += level;
				startPos++;
				chNext = styler[startPos];
				styleNext = styler.StyleIndexAt(startPos);
			}
		} break;

		case SCE_KOTLIN_RAWSTRING:
			if (style != stylePrev) {
				levelNext++;
			}
			if (style != styleNext) {
				levelNext--;
			}
			break;

		case SCE_KOTLIN_OPERATOR:
		case SCE_KOTLIN_OPERATOR2:
			if (ch == '{' || ch == '[' || ch == '(') {
				levelNext++;
			} else if (ch == '}' || ch == ']' || ch == ')') {
				levelNext--;
			}
			break;
		}

		if (visibleChars == 0 && !IsSpaceEquiv(style)) {
			++visibleChars;
		}
		if (startPos == lineStartNext) {
			const FoldLineState foldNext(styler.GetLineState(lineCurrent + 1));
			levelNext = sci::max(levelNext, SC_FOLDLEVELBASE);
			if (foldCurrent.lineComment) {
				levelNext += foldNext.lineComment - foldPrev.lineComment;
			} else if (foldCurrent.packageImport) {
				levelNext += foldNext.packageImport - foldPrev.packageImport;
			} else if (visibleChars) {
				const Sci_PositionU bracePos = CheckBraceOnNextLine(styler, lineCurrent, SCE_KOTLIN_OPERATOR, SCE_KOTLIN_TASKMARKER);
				if (bracePos) {
					levelNext++;
					startPos = bracePos + 1; // skip the brace
					style = SCE_KOTLIN_OPERATOR;
					chNext = styler[startPos];
					styleNext = styler.StyleIndexAt(startPos);
				}
			}

			const int levelUse = levelCurrent;
			int lev = levelUse | (levelNext << 16);
			if (levelUse < levelNext) {
				lev |= SC_FOLDLEVELHEADERFLAG;
			}
			styler.SetLevel(lineCurrent, lev);

			lineCurrent++;
			lineStartNext = styler.LineStart(lineCurrent + 1);
			levelCurrent = levelNext;
			foldPrev = foldCurrent;
			foldCurrent = foldNext;
			visibleChars = 0;
		}
	}
}

}

extern const LexerModule lmKotlin(SCLEX_KOTLIN, ColouriseKotlinDoc, "kotlin", FoldKotlinDoc);
